# coding=utf-8
import keras
from keras.layers import Input, Conv2D, BatchNormalization, MaxPooling2D
from keras.layers import Reshape, Dense, Dropout, Lambda, Activation
from keras.optimizers import Adam
from keras import backend as K
from keras.models import Model
from keras.utils import multi_gpu_model
import tensorflow as tf


def am_hparams():
    params = tf.contrib.training.HParams(
        # vocab
        vocab_size=None,
        lr=0.0008,
        gpu_nums=1,
        AUDIO_LENGTH=None,
        is_training=True)
    return params


# =============================搭建模型====================================
class Am():
    """docstring for Amodel."""

    def __init__(self, args):
        self.vocab_size = args.vocab_size
        self.gpu_nums = args.gpu_nums
        self.lr = args.lr
        self.is_training = args.is_training
        self.AUDIO_LENGTH = args.AUDIO_LENGTH
        self.AUDIO_FEATURE_LENGTH = 200  # 特征长度
        self._model_init()
        if self.is_training:
            self._ctc_init()
            self.opt_init()
            print("模型完成【编译】...")

    def _model_init(self):
        """
        		定义DFRestCNN模型，使用函数式模型 修改DFCNN 加入残差卷积
        		输入层：200维的特征值序列，一条语音数据的最大长度设为1600（大约16s）
        		隐藏层：卷积池化层，卷积核大小为3x3，池化窗口大小为2
        		隐藏层：全连接层
        		输出层：全连接层，神经元数量为self.vocab_size，使用softmax作为激活函数，
        		CTC层：使用CTC的loss作为损失函数，实现连接性时序多输出
        """
        self.inputs = Input(name='the_inputs', shape=(self.AUDIO_LENGTH, self.AUDIO_FEATURE_LENGTH, 1))
        layer_h1 = Conv2D(32, (3, 3), use_bias=False, activation='relu', padding='same',
                          kernel_initializer='he_normal')(self.inputs)  # 卷积层
        layer_h1 = BatchNormalization(mode=0, axis=-1)(layer_h1)
        # layer_h1 = Dropout(0.05)(layer_h1)
        layer_h2 = Conv2D(32, (3, 3), use_bias=True, activation='relu', padding='same', kernel_initializer='he_normal')(
            layer_h1)  # 卷积层
        layer_h2 = BatchNormalization(axis=-1)(layer_h2)
        layer_h3 = MaxPooling2D(pool_size=2, strides=None, padding="valid")(layer_h2)  # 池化层 800*100
        # layer_h3 = Dropout(0.2)(layer_h2) # 随机中断部分神经网络连接，防止过拟合
        # layer_h3 = Dropout(0.05)(layer_h3) # 随机中断部分神经网络连接，防止过拟合
        # layer_h3 = BatchNormalization(axis=1)(layer_h3)
        layer_h4 = Conv2D(64, (3, 3), use_bias=True, activation='relu', padding='same', kernel_initializer='he_normal')(
            layer_h3)  # 卷积层
        layer_h4 = BatchNormalization(axis=-1)(layer_h4)
        layer_h5 = Conv2D(64, (3, 3), use_bias=True, activation='relu', padding='same', kernel_initializer='he_normal')(
            layer_h4)  # 卷积层
        layer_h5 = BatchNormalization(axis=-1)(layer_h5)
        layer_h6 = MaxPooling2D(pool_size=2, strides=None, padding="valid")(layer_h5)  # 池化层 400*50
        # layer_h6 = Dropout(0.1)(layer_h6)
        layer_h7 = Conv2D(128, (3, 3), use_bias=True, activation='relu', padding='same',
                          kernel_initializer='he_normal')(layer_h6)  # 卷积层
        layer_h7 = BatchNormalization(axis=-1)(layer_h7)
        layer_h8 = Conv2D(128, (3, 3), use_bias=True, activation='relu', padding='same',
                          kernel_initializer='he_normal')(layer_h7)  # 卷积层
        layer_h8 = BatchNormalization(axis=-1)(layer_h8)
        layer_h9 = MaxPooling2D(pool_size=2, strides=None, padding="valid")(layer_h8)  # 池化层 200*25

        # layer_h9 = Dropout(0.15)(layer_h9)
        # layer_h9 = BatchNormalization(axis=-1)(layer_h9)
        layer_h10 = Conv2D(128, (3, 3), use_bias=True, activation='relu', padding='same',
                           kernel_initializer='he_normal')(layer_h9)  # 卷积层
        layer_h10 = BatchNormalization(axis=-1)(layer_h10)
        layer_h11 = Conv2D(128, (3, 3), use_bias=True, activation='relu', padding='same',
                           kernel_initializer='he_normal')(layer_h10)  # 卷积层
        layer_h11 = BatchNormalization(axis=-1)(layer_h11)
        layer_h12 = MaxPooling2D(pool_size=1, strides=None, padding="valid")(layer_h11)  # 池化层 200*25

        # layer_h12 = BatchNormalization(axis=-1)(layer_h12)
        layer_h13 = Conv2D(128, (3, 3), use_bias=True, activation='relu', padding='same',
                           kernel_initializer='he_normal')(layer_h12)  # 卷积层
        layer_h13 = BatchNormalization(axis=-1)(layer_h13)
        layer_h14 = Conv2D(128, (3, 3), use_bias=True, activation='relu', padding='same',
                           kernel_initializer='he_normal')(layer_h13)  # 卷积层
        layer_h14 = BatchNormalization(axis=-1)(layer_h14)
        layer_h15 = MaxPooling2D(pool_size=1, strides=None, padding="valid")(layer_h14)  # 池化层 None*200*25*128

        # test=Model(inputs = input_data, outputs = layer_h12)
        # test.summary()

        layer_h16 = Reshape((-1, int(self.AUDIO_FEATURE_LENGTH*128/8)))(layer_h15)  # Reshape层
        # layer_h5 = LSTM(256, activation='relu', use_bias=True, return_sequences=True)(layer_h4) # LSTM层
        # layer_h6 = Dropout(0.2)(layer_h5) # 随机中断部分神经网络连接，防止过拟合
        layer_h16 = Dropout(0.3)(layer_h16)
        # layer_h17 = BatchNormalization(axis=-1)(layer_h16)
        layer_h17 = Dense(128, activation="relu", use_bias=True, kernel_initializer='he_normal')(layer_h16)  # 全连接层
        layer_h17 = BatchNormalization(axis=-1)(layer_h17)
        layer_h17 = Dropout(0.3)(layer_h17)

        layer_h18 = Dense(self.vocab_size, use_bias=True, kernel_initializer='he_normal')(layer_h17)  # 全连接层
        layer_h18 = BatchNormalization(axis=-1)(layer_h18)
        self.outputs = Activation('softmax', name='Activation0')(layer_h18)
        self.model = Model(inputs=self.inputs, outputs=self.outputs)
        self.model.summary()

    def _ctc_init(self):
        self.labels = Input(name='the_labels', shape=[None], dtype='float32')
        self.input_length = Input(name='input_length', shape=[1], dtype='int64')
        self.label_length = Input(name='label_length', shape=[1], dtype='int64')
        self.loss_out = Lambda(ctc_lambda, output_shape=(1,), name='ctc') \
            ([self.labels, self.outputs, self.input_length, self.label_length])
        self.ctc_model = Model(inputs=[self.labels, self.inputs,
                                       self.input_length, self.label_length], outputs=self.loss_out)

    def opt_init(self):
        opt = Adam(lr=self.lr, beta_1=0.9, beta_2=0.999, decay=0.01, epsilon=10e-8)
        if self.gpu_nums > 1:
            self.ctc_model = multi_gpu_model(self.ctc_model, gpus=self.gpu_nums)
        self.ctc_model.compile(loss={'ctc': lambda y_true, output: output}, optimizer=opt)


# ============================模型组件=================================
def conv2d(size):
    return Conv2D(size, (3, 3), use_bias=True, activation='relu',
                  padding='same', kernel_initializer='he_normal')


def norm(x):
    return BatchNormalization(axis=-1)(x)


def maxpool(x):
    return MaxPooling2D(pool_size=(2, 2), strides=None, padding="valid")(x)


def dense(units, activation="relu"):
    return Dense(units, activation=activation, use_bias=True,
                 kernel_initializer='he_normal')


# x.shape=(none, none, none)
# output.shape = (1/2, 1/2, 1/2)
def cnn_cell(size, x, pool=True):
    x = norm(conv2d(size)(x))
    x = norm(conv2d(size)(x))
    if pool:
        x = maxpool(x)
    return x


def ctc_lambda(args):
    labels, y_pred, input_length, label_length = args
    y_pred = y_pred[:, :, :]
    return K.ctc_batch_cost(labels, y_pred, input_length, label_length)
